from ..common import VolcengineClient
from ..util import convert_bytes_to_image

models = ["通用2.1", "通用2.0Pro"]
req_schedule_confs = ["general_v20_9B_pe", "general_v20_9B_rephraser"]

model_configs = {
    "通用2.1": {
        "req_key": "high_aes_general_v21_L",
        "model_version": "general_v2.1_L",
        "llm_seed": -1,
        "seed": -1,
        "use_pre_llm": True,
        "use_sr": True,
        "sr_seed": -1,
        "sr_strength": 0.4,
        "sr_scale": 3.5,
        "sr_steps": 20,
        "is_only_sr": False,
    },
    "通用2.0Pro": {
        "req_key": "high_aes_general_v20_L",
        "model_version": "general_v2.0_L",
        "seed": -1,
        "use_sr": True,
    }
}


class DoubaoConfig:
    """
    豆包配置类
    """
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "load"
    CATEGORY = "DoubaoConfig"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "info": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "{}",
                    "lazy": True
                }),
            },
        }


class DouBao:
    RETURN_TYPES = ("IMAGE",)
    # RETURN_NAMES = ("image_output_name",)

    FUNCTION = "image_creator"

    # OUTPUT_NODE = False

    CATEGORY = "DouBao"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "info": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "{}",
                    "lazy": True
                }),
                "model": (models, {"default": models[0]}),
                "req_schedule_conf": (req_schedule_confs, {"default": req_schedule_confs[0]}),
                "prompt": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "",
                    "lazy": False
                }),
                "negative_prompt": ("STRING", {
                    "multiline": True,  # True if you want the field to look like the one on the ClipTextEncode node
                    "default": "",
                    "lazy": False
                }),
                "scale": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1000, "step": 0.05}),
                "steps": ("INT", {"default": 80, "min": 0, "max": 200, "step": 1}),
                "width": ("INT", {"default": 512, "min": 0, "max": 10000, "step": 1}),
                "height": ("INT", {"default": 512, "min": 0, "max": 10000, "step": 1}),
            },
        }

    def check_lazy_status(self, info, model, req_schedule_conf, prompt, negative_prompt, scale, steps, width, height):
        """
            Return a list of input names that need to be evaluated.

            This function will be called if there are any lazy inputs which have not yet been
            evaluated. As long as you return at least one field which has not yet been evaluated
            (and more exist), this function will be called again once the value of the requested
            field is available.

            Any evaluated inputs will be passed as arguments to this function. Any unevaluated
            inputs will have the value None.
        """
        pass

    @staticmethod
    def image_creator(info, model, req_schedule_conf, prompt, negative_prompt, scale, steps, width, height):
        image = None
        client = VolcengineClient(volcengine_info=info)
        operate_result = client.image_creator(schedule_conf=req_schedule_conf, steps=steps, width=width,
                                              height=height, prompt=prompt, negative_prompt=negative_prompt,
                                              scale=scale, config=model_configs.get(model, {}))
        if operate_result.code == 1:
            pics = operate_result.data
            image_bytes = pics[0]
            image = convert_bytes_to_image(image_bytes)
        return (image,)
